import math
import numpy as np
from scipy.spatial.transform import Rotation, Slerp
import torch
from scipy.spatial import ConvexHull

def quaternion_angle_difference(q1, q2):
    # 确保输入是四元数 (w, x, y, z)
    q1 = np.array(q1, dtype=float)
    q2 = np.array(q2, dtype=float)
    
    # 归一化四元数（确保它们是单位四元数）
    q1_norm = np.linalg.norm(q1)
    q2_norm = np.linalg.norm(q2)
    q1 = q1 / q1_norm
    q2 = q2 / q2_norm
    
    # 计算点积
    dot_product = np.dot(q1, q2)
    
    # 计算 cos(theta/2)，因为四元数表示的是旋转的一半角
    cos_half_theta = dot_product
    
    # 确保结果在 [-1, 1] 范围内，以避免数值误差导致的错误
    cos_half_theta = np.clip(cos_half_theta, -1.0, 1.0)
    
    # 计算夹角的一半（弧度）
    half_theta = np.arccos(cos_half_theta)
    
    # 计算完整的夹角（弧度）
    theta = 2 * half_theta
    
    # 将结果转换为度数（可选）
    theta_degrees = np.degrees(theta)
    
    return theta, theta_degrees

def quaternion_to_rotation_matrix(q):
    """
    Convert a quaternion to a rotation matrix.
    
    Parameters:
    q (list or numpy array): A quaternion [qx, qy, qz, qw]
    
    Returns:
    numpy array: A 3x3 rotation matrix
    """
    q = np.array(q, dtype=np.float64)
    q = q / np.linalg.norm(q)
    qx, qy, qz, qw = q
    
    R = np.array([
        [1 - 2*(qy**2 + qz**2),     2*(qx*qy - qz*qw),     2*(qx*qz + qy*qw)],
        [    2*(qx*qy + qz*qw), 1 - 2*(qx**2 + qz**2),     2*(qy*qz - qx*qw)],
        [    2*(qx*qz - qy*qw),     2*(qy*qz + qx*qw), 1 - 2*(qx**2 + qy**2)]
    ], dtype=np.float64)
    
    return R

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

# TODO: 根据时间插值计算位姿 pose[x, y, z, qx, qy, qz, qw]
def pose_interpolation(pose1, pose2, t1, t2, t):
    if t2 < t1:
        t2, t1 = t1, t2
        pose2, pose1 = pose1, pose2
    
    if t < t1 or t > t2:
        raise ValueError("t must be between t1 and t2")
    if abs(t2 - t1) < 1e-6:
        return pose1

    rate = (t - t1) / (t2 - t1)

    # 位置线性插值
    trans1 = np.array(pose1[:3])
    trans2 = np.array(pose2[:3])
    trans_interp = trans1 + rate * (trans2 - trans1)  # Lerp

    # 旋转插值
    q1 = Rotation.from_quat(pose1[3:])
    q2 = Rotation.from_quat(pose2[3:])
    slerp = Slerp([0, 1], Rotation.concatenate([q1, q2]))
    rot_interp = slerp(rate).as_quat()

    return np.concatenate([trans_interp, rot_interp])

def calculate_iou(box1, box2):
    if x1_min > x2_max or x2_min > x1_max or y1_min > y2_max or y2_min > y1_max:
        print("error box coordinates")
        return -1

    # 提取坐标
    x1_min, y1_min, x1_max, y1_max = box1
    x2_min, y2_min, x2_max, y2_max = box2

    # 计算交集区域坐标
    inter_x_min = max(x1_min, x2_min)
    inter_y_min = max(y1_min, y2_min)
    inter_x_max = min(x1_max, x2_max)
    inter_y_max = min(y1_max, y2_max)

    # 计算交集面积
    inter_w = max(0, inter_x_max - inter_x_min)
    inter_h = max(0, inter_y_max - inter_y_min)
    intersection = inter_w * inter_h

    # 计算各自面积及并集
    area1 = (x1_max - x1_min) * (y1_max - y1_min)
    area2 = (x2_max - x2_min) * (y2_max - y2_min)
    union = area1 + area2 - intersection

    # 处理除零错误
    if union == 0:
        return 0.0
    return intersection / union

# 默认box1的得分比box2高，用于判断box1是否包含box2
def calculate_cover(box1, box2, score1, score2):

    if score1 < score2:
        box1, box2 = box2, box1
        score1, score2 = score2, score1

    # 提取坐标
    x1_min, y1_min, x1_max, y1_max = box1
    x2_min, y2_min, x2_max, y2_max = box2

    if x1_min > x1_max or y1_min > y1_max or x2_min > x2_max or y2_min > y2_max:
        print("error box coordinates")
        return -1

    # 计算交集区域坐标
    inter_x_min = max(x1_min, x2_min)
    inter_y_min = max(y1_min, y2_min)
    inter_x_max = min(x1_max, x2_max)
    inter_y_max = min(y1_max, y2_max)

    # 计算交集面积
    inter_w = max(0, inter_x_max - inter_x_min)
    inter_h = max(0, inter_y_max - inter_y_min)
    intersection = inter_w * inter_h

    # 计算各自面积
    area1 = (x1_max - x1_min) * (y1_max - y1_min)
    area2 = (x2_max - x2_min) * (y2_max - y2_min)
    if area1 == 0 or area2 == 0:
        print("box area is 0")
        return -1

    # 计算覆盖率
    return max(intersection / area2, intersection / area1)

def compute_3d_iou_aabb(box1, box2):
    """
    计算两个三维轴对齐边界框的IoU
    参数格式: box1和box2均为元组或列表, 表示两个对角顶点的坐标:
        box = (x_min, y_min, z_min, x_max, y_max, z_max)
    """
    # 提取坐标轴范围
    x1_min, y1_min, z1_min, x1_max, y1_max, z1_max = box1
    x2_min, y2_min, z2_min, x2_max, y2_max, z2_max = box2
    
    # 计算交集区域的坐标边界
    xA = max(x1_min, x2_min)
    yA = max(y1_min, y2_min)
    zA = max(z1_min, z2_min)
    xB = min(x1_max, x2_max)
    yB = min(y1_max, y2_max)
    zB = min(z1_max, z2_max)
    
    # 计算交集体积
    inter_length = max(0, xB - xA)
    inter_width = max(0, yB - yA)
    inter_height = max(0, zB - zA)
    inter_volume = inter_length * inter_width * inter_height
    
    # 计算各自体积
    vol1 = (x1_max - x1_min) * (y1_max - y1_min) * (z1_max - z1_min)
    vol2 = (x2_max - x2_min) * (y2_max - y2_min) * (z2_max - z2_min)
    union_volume = vol1 + vol2 - inter_volume
    
    # 计算IoU
    iou = inter_volume / union_volume if union_volume != 0 else 0.0
    return iou

def compute_obb_iou(obb1, obb2):
    """
    计算两个OBB的IoU，使用PyTorch张量加速
    参数：
        obb1/obb2: Open3D的OBB对象，包含center, extent, R属性
    返回：
        iou (float): 交并比 [0,1]
    """
    # 将OBB参数转换为PyTorch张量并转移到GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    center1 = torch.tensor(obb1.center, dtype=torch.float32).to(device)
    extent1 = torch.tensor(obb1.extent, dtype=torch.float32).to(device)
    R1 = torch.tensor(obb1.R, dtype=torch.float32).to(device)
    
    center2 = torch.tensor(obb2.center, dtype=torch.float32).to(device)
    extent2 = torch.tensor(obb2.extent, dtype=torch.float32).to(device)
    R2 = torch.tensor(obb2.R, dtype=torch.float32).to(device)
    
    # 计算OBB的8个顶点（使用PyTorch矩阵运算）
    def get_box_points(center, extent, R):
        # 生成立方体顶点在局部坐标系下的坐标
        x = torch.tensor([1, -1, 1, -1, 1, -1, 1, -1], dtype=torch.float32)
        y = torch.tensor([1, 1, -1, -1, 1, 1, -1, -1], dtype=torch.float32)
        z = torch.tensor([1, 1, 1, 1, -1, -1, -1, -1], dtype=torch.float32)
        points = torch.stack([x, y, z], dim=1).view(8, 3).to(device)
        # 应用旋转矩阵和中心平移
        rotated = torch.matmul(points, R.T) * extent.view(1, 3)
        return rotated + center.view(1, 3)
    
    # 获取顶点并转换为张量
    points1 = get_box_points(center1, extent1, R1).detach().cpu().numpy()
    points2 = get_box_points(center2, extent2, R2).detach().cpu().numpy()
    
    # 使用scipy计算凸包交集体积（需安装scipy）
    convex_hull1 = ConvexHull(points1)
    convex_hull2 = ConvexHull(points2)
    
    # 计算交集体积（简化版，实际需三维布尔运算）
    # 这里使用凸包的并集减去各自体积作为近似
    volume1 = convex_hull1.volume
    volume2 = convex_hull2.volume
    union_volume = volume1 + volume2 - convex_hull1.intersection(convex_hull2).volume
    
    # 计算IoU
    inter_volume = convex_hull1.intersection(convex_hull2).volume
    iou = inter_volume / union_volume if union_volume != 0 else 0.0
    return iou

def compute_aabb_iou(aabb1, aabb2):
    """
    计算两个AABB的IoU
    参数：
        aabb1/2: Open3D的AxisAlignedBoundingBox对象
    返回：
        iou (float): 交并比 [0,1]
    """
    # 提取AABB的最小/最大坐标
    min1, max1 = np.asarray(aabb1.min_bound), np.asarray(aabb1.max_bound)
    min2, max2 = np.asarray(aabb2.min_bound), np.asarray(aabb2.max_bound)
    
    # 计算交集区域边界
    inter_min = np.maximum(min1, min2)
    inter_max = np.minimum(max1, max2)
    
    # 计算交集各轴长度（若无交集则为0）
    inter_dims = np.maximum(inter_max - inter_min, 0)
    inter_volume = np.prod(inter_dims)
    
    # 计算各自体积
    vol1 = np.prod(max1 - min1)
    vol2 = np.prod(max2 - min2)
    union_volume = vol1 + vol2 - inter_volume
    
    # 计算IoU
    iou = inter_volume / union_volume if union_volume != 0 else 0.0
    return iou